import numpy as np
import tensorflow as tf
import copy
import time  # For potential timing print "耗时"

# Import necessary functions from your adaptive_attack.py
# Make sure adaptive_attack.py is in the same directory or accessible in PYTHONPATH
from adaptive_attack import craft_heuristic_malicious_updates, flatten_weights, get_shapes_and_dtypes


class Server:
    def __init__(self, model_factory, select_client, iteration, alpha, beta):
        self._model = model_factory()
        self.nselect_client = select_client
        self.global_itera = iteration
        self.alpha = alpha
        self.beta = beta  # <--- 确保这一行已添加/修正

        self._model.compile(
            optimizer=tf.keras.optimizers.SGD(learning_rate=self.beta),  # 使用 self.beta
            loss=tf.keras.losses.CategoricalCrossentropy(),
            metrics=['accuracy']
        )

    def train_self(self, root_client, expr_basename, x_test, y_test):
        """Trains the model using only the root_client's data (centralized training on server's trusted data)."""
        for r in range(0, self.global_itera):
            server_weights = self._model.get_weights()
            root_delta = root_client.train(copy.deepcopy(server_weights))  # Get update from root client

            # Update global model: W_new = W_old + alpha * delta
            new_server_weights = [old_w + self.alpha * delta_w for delta_w, old_w in zip(server_weights, root_delta)]
            self._model.set_weights(new_server_weights)

            # print(f'{expr_basename} round={r + 1}/{self.global_itera}', end='') # 保留下面一行合并的评估信息
            loss, acc = self._model.evaluate(x_test, y_test, verbose=0, batch_size=32)
            print(f'{expr_basename} round={r + 1}/{self.global_itera} loss: {loss:.4f} - accuracy: {acc:.2%}') # 修改了这里的打印格式以匹配最终评估输出

    def evaluate(self, x_test, y_test, expr_basename):
        loss, acc = self._model.evaluate(x_test, y_test, verbose=0, batch_size=32)
        print(f'{expr_basename} loss: {loss:.4f} - accuracy: {acc:.2%}')

    def train(self, clients, root_client, expr_basename, x_test, y_test):
        """
        Original Federated Learning training with FLTrust aggregation.
        Assumes all participating clients are benign.
        """
        for r in range(0, self.global_itera):
            # print(f'{expr_basename} round={r + 1}/{self.global_itera}') # 保留下面一行合并的评估信息
            server_weights = self._model.get_weights()  # W_t

            selected_clients = clients if self.nselect_client == len(clients) \
                else np.random.choice(clients, self.nselect_client, replace=False)

            client_deltas = []
            for i, client_obj in enumerate(selected_clients):
                delta = client_obj.train(copy.deepcopy(server_weights))
                client_deltas.append(delta)

            g0 = root_client.train(copy.deepcopy(server_weights))

            aggregated_delta = [np.zeros_like(w) for w in server_weights]
            g0_flat = flatten_weights(g0)
            g0_norm = np.linalg.norm(g0_flat)
            total_TS = 0.0
            temp_weighted_deltas = []

            for delta_i in client_deltas:
                delta_i_flat = flatten_weights(delta_i)
                delta_i_norm = np.linalg.norm(delta_i_flat)
                cos_sim = 0.0
                if delta_i_norm > 1e-9 and g0_norm > 1e-9:
                    cos_sim = np.dot(delta_i_flat, g0_flat) / (delta_i_norm * g0_norm)
                TS_i = np.maximum(0, cos_sim)
                total_TS += TS_i
                norm_ratio = 0.0
                if delta_i_norm > 1e-9: # 确保 g0_norm 在之前已检查 > 1e-9
                    norm_ratio = g0_norm / delta_i_norm if g0_norm > 1e-9 else 0.0
                term_for_sum = [TS_i * norm_ratio * w for w in delta_i]
                temp_weighted_deltas.append(term_for_sum)

            if total_TS > 1e-9:
                for term in temp_weighted_deltas:
                    for k_layer in range(len(aggregated_delta)):
                        aggregated_delta[k_layer] += term[k_layer]
                for k_layer in range(len(aggregated_delta)):
                    aggregated_delta[k_layer] /= total_TS
            # else:
                # print("    [Warning] Total Trust Score is zero in original train. Aggregated update will be zero.") # 已注释

            new_global_weights = [
                sw_k + self.alpha * agg_d_k
                for sw_k, agg_d_k in zip(server_weights, aggregated_delta)
            ]
            self._model.set_weights(new_global_weights)

            loss, acc = self._model.evaluate(x_test, y_test, verbose=0, batch_size=32)
            # 保持和 train_adaptive_attack 一致的输出格式，包含 expr_basename 和 round
            print(f'{expr_basename} round={r + 1}/{self.global_itera} 评估 - loss: {loss:.4f}, accuracy: {acc:.2%}')

    def train_smarter_adaptive_attack(self, clients, root_client, expr_basename, x_test, y_test,
                                      initial_benign_count=50,
                                      malicious_config={'poor_benign': 20, 'good_benign': 50},
                                      benign_effectiveness_threshold=0.5,
                                      cos_candidates=np.arange(0, 0.45, 0.05)  # 确保这是个列表或ndarray
                                      ):
        print(f"Starting smarter adaptive attack training: {expr_basename}")
        total_clients = len(clients)

        if not isinstance(cos_candidates, (list, np.ndarray)):
            print("[Error] cos_candidates should be a list or numpy array. Using default.")
            cos_candidates = np.arange(0, 0.45, 0.05)
        if len(cos_candidates) == 0:
            print("[Error] cos_candidates is empty. Using default.")
            cos_candidates = np.arange(0, 0.45, 0.05)

        for r in range(self.global_itera):
            print(f'{expr_basename} round={r + 1}/{self.global_itera}')
            server_weights_r = self._model.get_weights()

            g0 = root_client.train(copy.deepcopy(server_weights_r))
            g0_flat = flatten_weights(g0)
            g0_norm = np.linalg.norm(g0_flat)

            if g0_norm < 1e-9:
                print("    [Warning] g0_norm is close to zero. Applying g0 directly for server update.")
                # 如果没有有效的g0，则服务器模型仅通过alpha * g0更新（如果g0本身有意义）
                # 或者，如果g0也接近零，则模型可能不会有效更新。
                # 这里的策略是至少尝试应用g0。
                final_aggregated_delta = g0  # 使用g0作为本轮的聚合结果
                new_global_weights = [sw_k + self.alpha * agg_d_k for sw_k, agg_d_k in
                                      zip(server_weights_r, final_aggregated_delta)]
                self._model.set_weights(new_global_weights)
                loss, acc = self._model.evaluate(x_test, y_test, verbose=0, batch_size=32)
                print(f'    评估 (g0_norm near zero, used g0) - loss: {loss:.4f}, accuracy: {acc:.2%}')
                continue  # 跳过本轮剩余的攻击逻辑

            shapes_dtypes = get_shapes_and_dtypes(g0)

            # 为本轮的模拟评估创建并编译一次评估模型
            eval_model_for_simulation = tf.keras.models.clone_model(self._model)
            eval_model_for_simulation.compile(
                optimizer=tf.keras.optimizers.SGD(learning_rate=self.beta),
                loss=tf.keras.losses.CategoricalCrossentropy(),
                metrics=['accuracy']
            )

            # 1. 动态客户端分配
            benign_deltas_this_round = []
            actual_num_malicious_this_round = 0

            if total_clients < initial_benign_count:
                print(
                    f"    [Warning] Not enough total clients ({total_clients}) for initial_benign_count ({initial_benign_count}). Treating all as benign.")
                for client_obj in clients:
                    delta = client_obj.train(copy.deepcopy(server_weights_r))
                    benign_deltas_this_round.append(delta)
                # actual_num_malicious_this_round  остается 0
            else:
                shuffled_clients = list(clients)  # 创建可修改的副本
                np.random.shuffle(shuffled_clients)

                initial_benign_clients_this_round = shuffled_clients[:initial_benign_count]
                # 构造一个不包含 initial_benign_clients_this_round 的列表
                # 使用 set 操作来确保 non_initial_clients 是正确的
                initial_benign_set = set(initial_benign_clients_this_round)
                potential_remaining_clients = [c for c in shuffled_clients if c not in initial_benign_set]

                initial_benign_deltas_collected = []  # 用于收集初始良性更新
                sum_cos_sim_initial_benign = 0.0
                valid_initial_benign_updates_count = 0

                for client_obj in initial_benign_clients_this_round:
                    delta = client_obj.train(copy.deepcopy(server_weights_r))
                    initial_benign_deltas_collected.append(delta)
                    delta_flat = flatten_weights(delta)
                    delta_norm = np.linalg.norm(delta_flat)
                    if delta_norm > 1e-9:
                        cos_sim = np.dot(delta_flat, g0_flat) / (delta_norm * g0_norm)
                        sum_cos_sim_initial_benign += cos_sim
                        valid_initial_benign_updates_count += 1

                avg_cos_sim_initial_benign = (
                            sum_cos_sim_initial_benign / valid_initial_benign_updates_count) if valid_initial_benign_updates_count > 0 else 0.0
                benign_deltas_this_round.extend(initial_benign_deltas_collected)  # 将初始良性更新加入本轮列表

                if avg_cos_sim_initial_benign < benign_effectiveness_threshold:
                    print(f"    Initial benign updates deemed 'poor' (avg_cos_sim={avg_cos_sim_initial_benign:.3f}).")
                    actual_num_malicious_this_round = malicious_config['poor_benign']
                    actual_num_malicious_this_round = min(actual_num_malicious_this_round,
                                                          len(potential_remaining_clients))  # 不能超过可用数量

                    num_additional_benign = len(potential_remaining_clients) - actual_num_malicious_this_round
                    if num_additional_benign > 0:
                        # 从 potential_remaining_clients 中选择额外的良性客户端
                        additional_benign_clients_this_round = potential_remaining_clients[:num_additional_benign]
                        for client_obj in additional_benign_clients_this_round:
                            delta = client_obj.train(copy.deepcopy(server_weights_r))
                            benign_deltas_this_round.append(delta)
                else:
                    print(f"    Initial benign updates deemed 'good' (avg_cos_sim={avg_cos_sim_initial_benign:.3f}).")
                    actual_num_malicious_this_round = malicious_config['good_benign']
                    actual_num_malicious_this_round = min(actual_num_malicious_this_round,
                                                          len(potential_remaining_clients))

            print(
                f"    This round: Benign Deltas={len(benign_deltas_this_round)}, Malicious Updates to craft={actual_num_malicious_this_round}")

            # --- [攻击者视角] 模拟以确定最佳 cos ---
            best_cos_for_this_round = cos_candidates[0]
            min_simulated_accuracy = float('inf')

            if actual_num_malicious_this_round > 0 and len(benign_deltas_this_round) > 0:
                for cos_candidate in cos_candidates:
                    hypothetical_malicious_deltas = craft_heuristic_malicious_updates(
                        g0, benign_deltas_this_round, actual_num_malicious_this_round,
                        shapes_dtypes, cos_candidate
                    )

                    sim_aggregated_delta = [np.zeros_like(w) for w in server_weights_r]
                    sim_temp_weighted_contribs = []
                    sim_total_TS = 0.0

                    all_sim_deltas = hypothetical_malicious_deltas + benign_deltas_this_round
                    for delta_obj in all_sim_deltas:
                        delta_obj_flat = flatten_weights(delta_obj)
                        delta_obj_norm = np.linalg.norm(delta_obj_flat)
                        cos_sim_g0 = 0.0
                        if delta_obj_norm > 1e-9:
                            cos_sim_g0 = np.dot(delta_obj_flat, g0_flat) / (delta_obj_norm * g0_norm)
                        TS_i = np.maximum(0, cos_sim_g0)
                        sim_total_TS += TS_i
                        norm_ratio = g0_norm / (delta_obj_norm + 1e-9)
                        weighted_contrib = [TS_i * norm_ratio * w for w in delta_obj]
                        sim_temp_weighted_contribs.append(weighted_contrib)

                    if sim_total_TS > 1e-9:
                        for contrib in sim_temp_weighted_contribs:
                            for k_layer in range(len(sim_aggregated_delta)):
                                sim_aggregated_delta[k_layer] += contrib[k_layer]
                        for k_layer in range(len(sim_aggregated_delta)):
                            sim_aggregated_delta[k_layer] /= sim_total_TS

                    candidate_global_weights = [sw_k + self.alpha * agg_d_k for sw_k, agg_d_k in
                                                zip(server_weights_r, sim_aggregated_delta)]

                    eval_model_for_simulation.set_weights(candidate_global_weights)
                    _, sim_accuracy = eval_model_for_simulation.evaluate(x_test, y_test, verbose=0, batch_size=32)

                    if sim_accuracy < min_simulated_accuracy:
                        min_simulated_accuracy = sim_accuracy
                        best_cos_for_this_round = cos_candidate

                print(
                    f"    Attacker selected cos={best_cos_for_this_round:.2f} for this round (sim_acc={min_simulated_accuracy:.2%}).")
            elif actual_num_malicious_this_round == 0:
                print(f"    No malicious updates to craft this round. Skipping cos selection.")
            else:  # 有恶意数量但没有良性上下文
                print(
                    f"    Malicious updates configured, but no benign deltas for context. Cannot craft heuristic updates. Skipping attack.")
                actual_num_malicious_this_round = 0  # 重置，因为无法生成

            # --- 为本轮生成实际的恶意更新 ---
            actual_malicious_updates_for_server = []
            if actual_num_malicious_this_round > 0 and len(benign_deltas_this_round) > 0:
                actual_malicious_updates_for_server = craft_heuristic_malicious_updates(
                    g0, benign_deltas_this_round, actual_num_malicious_this_round,
                    shapes_dtypes, best_cos_for_this_round
                )

            # --- 服务器的常规FLTrust聚合 ---
            final_aggregated_delta = [np.zeros_like(w) for w in server_weights_r]

            all_final_deltas_to_aggregate = actual_malicious_updates_for_server + benign_deltas_this_round
            if not all_final_deltas_to_aggregate:  # 如果没有任何更新
                print(
                    "    [Warning] No client updates (benign or malicious) available for aggregation. Using g0 for update.")
                final_aggregated_delta = copy.deepcopy(g0)  # 使用g0的副本
            else:
                final_temp_weighted_contribs = []
                final_total_TS = 0.0
                for delta_obj in all_final_deltas_to_aggregate:
                    delta_obj_flat = flatten_weights(delta_obj)
                    delta_obj_norm = np.linalg.norm(delta_obj_flat)
                    cos_sim_g0 = 0.0
                    if delta_obj_norm > 1e-9:
                        cos_sim_g0 = np.dot(delta_obj_flat, g0_flat) / (delta_obj_norm * g0_norm)
                    TS_i = np.maximum(0, cos_sim_g0)
                    final_total_TS += TS_i
                    norm_ratio = g0_norm / (delta_obj_norm + 1e-9)
                    weighted_contrib = [TS_i * norm_ratio * w for w in delta_obj]
                    final_temp_weighted_contribs.append(weighted_contrib)

                if final_total_TS > 1e-9:
                    for contrib in final_temp_weighted_contribs:
                        for k_layer in range(len(final_aggregated_delta)):
                            final_aggregated_delta[k_layer] += contrib[k_layer]
                    for k_layer in range(len(final_aggregated_delta)):
                        final_aggregated_delta[k_layer] /= final_total_TS
                else:  # 没有可信的更新
                    print("    [Warning] Total Trust Score (final_total_TS) is zero. Aggregated update will be zero.")
                    # final_aggregated_delta 保持为零向量

            # --- 更新服务器的全局模型 ---
            new_global_weights = [sw_k + self.alpha * agg_d_k for sw_k, agg_d_k in
                                  zip(server_weights_r, final_aggregated_delta)]
            self._model.set_weights(new_global_weights)

            # --- 评估并记录本轮结果 ---
            loss, acc = self._model.evaluate(x_test, y_test, verbose=0, batch_size=32)
            print(f'    评估 - loss: {loss:.4f}, accuracy: {acc:.2%}')

        # （可选）清理模拟模型
        if 'eval_model_for_simulation' in locals() and isinstance(eval_model_for_simulation, tf.keras.Model):
            del eval_model_for_simulation
            # print("Cleaned up eval_model_for_simulation.")